5.5. AI/ML Experiments

What is an AI/ML experiment?

An AI/ML experiment is a systematic and iterative process for building robust machine learning models. It involves testing different algorithms, tuning hyperparameters, and using various datasets to discover the optimal configuration for a specific predictive task. Each experiment is a structured trial designed to measure the impact of changes on model performance, such as accuracy, efficiency, and reliability.

Why is experiment tracking essential in AI/ML?

In MLOps, the complexity and often non-deterministic nature of model development require a disciplined approach. Experiment tracking provides the necessary structure, much like a lab notebook for a scientist. Key benefits include:

Ensuring Reproducibility: Tracking guarantees that every aspect of an experiment—code, data, environment, and parameters—is recorded. This allows you and your team to reliably replicate and verify results.
Optimizing Hyperparameters: It provides a systematic way to test and compare different hyperparameter configurations, helping you pinpoint the settings that maximize model performance.
Organizing Your Work: By logging experiments and using tags, you can categorize runs by model type, dataset, or objective. This organization is crucial for managing complex projects and quickly retrieving past results.
Monitoring Performance: Tracking metrics during a run offers real-time insight into how adjustments affect model behavior, enabling faster, data-driven decisions.
Seamless Framework Integration: Modern tracking tools integrate with popular AI/ML frameworks, creating a unified and streamlined workflow across your entire toolchain.

Which experiment tracking solution should you use?

Numerous solutions are available for tracking AI/ML experiments. Major cloud platforms like Google Cloud (Vertex AI), Azure (Azure ML), and AWS (SageMaker) offer powerful, integrated MLOps capabilities. Specialized commercial tools like Weights & Biases, Comet, and ClearML also provide excellent features.

A note on vendor risk: this chapter used to recommend Neptune.ai here. OpenAI acquired it in December 2025 and it wound down its external service, taking its documentation with it. That is worth internalising before you build a workflow on a hosted tracker: an experiment store is where your history lives, so prefer one whose format you can export, and treat a vendor's continued existence as an assumption rather than a given. It is one of the reasons this course teaches MLflow, which you can run yourself.

For those starting out or preferring an open-source, framework-agnostic solution, MLflow is an outstanding choice. It is versatile, robust, and integrates with a wide array of ML libraries.

To install MLflow, run:

uv add mlflow

To verify the installation and start the MLflow UI server locally:

uv run mlflow doctor
uv run mlflow server --backend-store-uri=sqlite:///mlflow.db --artifacts-destination=./mlruns

In the MLOps Python Package, that command is wrapped as a task, so you never have to retype the flags:

mise run mlflow:serve

For a more permanent setup using Docker, you can use a docker-compose.yml file to launch the MLflow server on the same store:

# From the reference-project root: initializes the SQLite file before Compose binds it
mise run docker:compose

Read the maintained Compose configuration. Its database and artifact directory are local persistent state. A client using the HTTP server must set its tracking URI accordingly and use artifact locations that the server can access; a shared database alone does not make host paths readable inside a container.

Run docker compose up to start the service. For information on production-grade deployments, refer to the MLflow documentation.

Where should MLflow store your experiments?

MLflow splits storage into two halves, and they do not have to live in the same place:

The backend store holds the metadata: experiments, runs, parameters, metrics, tags, and every entry in the model registry. It can be a local directory (the file store) or any SQLAlchemy-compatible database (SQLite, PostgreSQL, MySQL).
The artifact store holds the files: serialized models, plots, data samples, and anything else you log as an artifact. It is a filesystem path or an object store such as S3 or GCS.

This course uses SQLite for the backend store and the local ./mlruns directory for artifacts:

sqlite:///mlflow.db   # metadata: experiments, runs, metrics, registered models
./mlruns              # artifact files: models, plots, datasets

Three reasons make this the right default, even on a laptop:

It is the shape you will run in production. A production MLflow deployment uses a SQLAlchemy backend, almost always PostgreSQL. Developing against SQLite means the same store type, the same schema, and the same query semantics. Moving up to PostgreSQL becomes a change of connection string (MLFLOW_TRACKING_URI), not a rewrite of how your code talks to MLflow.
The database supports tracking and registry metadata in one place. Current MLflow disables its legacy file store by default. Use a database for this course rather than relying on legacy filesystem metadata. The next section builds directly on top of this choice.
You are agreeing with upstream, not working around it. Since MLflow 3.14, sqlite:///mlflow.db is MLflow's own default tracking URI (DEFAULT_TRACKING_URI in mlflow/store/tracking/__init__.py). MLflow keeps a backward-compatibility check that falls back to ./mlruns only when it detects an existing file store on disk. Setting the URI explicitly documents the intent and removes the ambiguity; it does not fight the library.

!!! warning "The file store is a legacy path" Earlier versions of this course configured ./mlruns as the backend store. If you have an existing ./mlruns directory holding file-store runs, MLflow will silently keep using it instead of SQLite. Back up the existing store before migration and preserve its historical runs. Set the tracking and registry URIs explicitly for new runs. Switching to SQLite does not migrate file-store metadata automatically; keep the old store available until you have verified your migration.

How do you configure MLflow in a project?

To integrate MLflow, you first need to point it at the backend store. Set the tracking and registry URIs to your SQLite database, then define an experiment name to group related runs.

Enabling MLflow's autologging is highly recommended. It automatically captures metrics, parameters, and models from popular ML libraries without requiring explicit logging statements.

import mlflow

# Metadata goes to the SQLite database; artifact files land under ./mlruns
mlflow.set_tracking_uri("sqlite:///mlflow.db")
mlflow.set_registry_uri("sqlite:///mlflow.db")

# Set a name for the experiment
mlflow.set_experiment(experiment_name="Bike Sharing Demand Prediction")

# Enable autologging for automatic tracking
mlflow.autolog()

In the MLOps Python Package, the maintained MlflowService gives each job the same URI rules: explicit YAML values take precedence over MLFLOW_TRACKING_URI and MLFLOW_REGISTRY_URI. Tracking defaults to sqlite:///mlflow.db; the registry defaults to the tracking URI when its environment variable is absent. Experiment and registry names are configured explicitly in the service.

from bikes.io.services import MlflowService

service = MlflowService()
assert service.tracking_uri
assert service.registry_uri

Remember to add both mlflow.db and mlruns/ to your .gitignore: experiment metadata and model artifacts are outputs, not source code.

To start a new run, wrap your training code within an MLflow run context. This allows you to add descriptive metadata and enable system metric logging.

with mlflow.start_run(
    run_name="Forecast Model with Feature Engineering",
    description="Training a model with an enhanced feature set.",
    log_system_metrics=True,
) as run:
    # Your model training and evaluation code goes here
    print(f"MLflow Run ID: {run.info.run_id}")

What information can you track in an experiment?

While MLflow's autologging captures a wealth of information automatically, you can enhance it with manual logging for greater detail:

Parameters: Log key-value parameters with mlflow.log_param() or a dictionary of parameters with mlflow.log_params().
Metrics: Record single metrics over time (e.g., per epoch) with mlflow.log_metric() or a dictionary of metrics with mlflow.log_metrics().
Inputs: Log dataset details and context with mlflow.log_input(), including tags for better categorization.
Tags: Assign custom tags to a run for improved filtering and organization using mlflow.set_tag() or mlflow.set_tags().
Artifacts: Save output files, such as plots, model files, or data samples, with mlflow.log_artifact() for single files or mlflow.log_artifacts() for directories.

How can you compare experiments to find the best model?

Comparing experiments is essential for model selection. MLflow provides two powerful ways to do this: its web UI and its programmatic API.

Comparing via the MLflow Web UI

The MLflow UI offers an intuitive, visual way to compare runs.

Launch the MLflow Server: If it's not running, start it with mise run mlflow:serve.
Select Runs: Navigate to the experiment page, where all runs are listed. Use the checkboxes to select the runs you want to compare.
Click Compare: A "Compare" button will appear. Clicking it opens a detailed view that places the selected runs side-by-side.
Analyze Results: This view provides a comprehensive summary of parameters, metrics, and artifacts for each run. You can use it to identify which configurations yielded the best performance.

Comparing Programmatically

The following is a classification illustration that requires runs with a logged validation_accuracy. The reference bike project is a regression task; compare its same named metric on the same evaluation period and keep the score's sign convention consistent.

Programmatic comparison is ideal for automated analysis and custom reporting.

Query Runs: Use mlflow.search_runs() to fetch run data into a pandas DataFrame. You can filter by experiment, metrics, parameters, or tags.

import mlflow

# Fetch runs from specific experiments
experiment_ids = ["1", "2"]
runs_df = mlflow.search_runs(experiment_ids)

Filter and Sort: With the data in a DataFrame, you can use pandas to sort and filter the results to find the top-performing runs based on your criteria.

# Find the best run based on validation accuracy
best_run = runs_df.sort_values("metrics.validation_accuracy", ascending=False).iloc[0]
print(f"Best Run ID: {best_run.run_id}")

Visualize Comparisons: Use libraries like Matplotlib or Seaborn to create custom visualizations that make comparisons clear and intuitive.

import matplotlib.pyplot as plt

# Plot validation accuracy for the top 5 runs
top_5_runs = runs_df.sort_values("metrics.validation_accuracy", ascending=False).head(5)
plt.figure(figsize=(12, 7))
plt.bar(top_5_runs["run_id"].str[:7], top_5_runs["metrics.validation_accuracy"])
plt.title("Comparison of Validation Accuracy Across Top 5 Runs")
plt.xlabel("Run ID")
plt.ylabel("Validation Accuracy")
plt.show()

What are some best practices for experiment tracking?

To maximize the value of your experiments, adopt these practices:

Use Clear Naming Conventions: Give experiments and runs descriptive names to make them easily identifiable (e.g., PROD_Retraining_ResNet50 vs. test_run_1).
Align with Business Metrics: Ensure that the metrics you track are directly relevant to project goals and business outcomes.

Leverage Nested Runs: Use nested runs to organize complex experiments, such as hyperparameter tuning, where each child run explores a different parameter set.

with mlflow.start_run(run_name="Hyperparameter Search") as parent_run:
    params = [0.01, 0.02, 0.03]
    for p in params:
        with mlflow.start_run(nested=True, run_name=f"alpha_{p}") as child_run:
            mlflow.log_param("alpha", p)
            # ... training logic ...
            mlflow.log_metric("val_loss", val_loss)

Tag Extensively: Use tags to add metadata like the dataset version, model type, or evaluation status (e.g., dataset:v2, model:xgboost, status:production_candidate).
Track Progress Over Time: Log metrics at each step or epoch to visualize the learning process and diagnose issues like overfitting.
```
# Inside your training loop
mlflow.log_metric(key="train_loss", value=train_loss, step=epoch)
```
Register Promising Models: When a run produces a high-quality model, log it to the MLflow Model Registry to version it and prepare it for deployment.

How do you test code that tracks experiments?

A real backend store is not free, and your test suite is where you feel it. Every time MLflow opens a SQLite database that does not exist yet, it runs its Alembic schema migrations before the first row is written. That costs a few seconds — around seven in the MLOps Python Package — and it happens per database.

The naive fixture creates a fresh database for every test, so the suite pays that price on every single test. When the reference package moved from the file store to SQLite with a per-test database, its suite went from 34 seconds to 339 seconds: a tenfold regression caused entirely by re-running the same migrations hundreds of times.

The fix keeps the isolation and drops the cost. Build one migrated database in a session-scoped fixture, then copy the file for each test. A file copy takes about 0.07 seconds instead of seven, and each test still gets a private, empty store it can write to freely. With this fixture, the reference suite settles at 63 seconds — the honest price of testing against the store you actually ship.

@pytest.fixture(scope="session")
def mlflow_db_template(tmp_path_factory: pytest.TempPathFactory) -> str:
    """Return a migrated but empty MLflow database used as a template by every test.

    Creating an MLflow SQLite store runs its Alembic migrations, which costs seconds.
    Paying that once per session and copying the file per test keeps the isolation of a
    fresh database at the cost of a file copy.
    """
    path = tmp_path_factory.mktemp("mlflow") / "template.db"
    services.MlflowService(
        tracking_uri=f"sqlite:///{path}",
        registry_uri=f"sqlite:///{path}",
        experiment_name="Experiment-Template",
        registry_name="Registry-Template",
    ).start()
    return str(path)


@pytest.fixture(scope="function", autouse=True)
def mlflow_service(tmp_path: str, mlflow_db_template: str) -> T.Generator[services.MlflowService]:
    """Return and start the mlflow service."""
    # Each test gets its own SQLite file under tmp_path, so runs stay isolated the same
    # way the old per-test directories were, with the store the package actually ships.
    database = os.path.join(tmp_path, "mlflow.db")
    shutil.copyfile(mlflow_db_template, database)
    service = services.MlflowService(
        tracking_uri=f"sqlite:///{database}",
        registry_uri=f"sqlite:///{database}",
        experiment_name="Experiment-Testing",
        registry_name="Registry-Testing",
    )
    service.start()
    yield service
    service.stop()

The general lesson outlives MLflow: when a test fixture is slow, look for the expensive, deterministic setup step hiding inside it and pay for it once. A migrated schema, a compiled model, or a seeded dataset can almost always be built at session scope and cloned per test.

How does experiment tracking fit into the MLOps lifecycle?

Experiment tracking is a cornerstone of the MLOps lifecycle, bridging the gap between development and production.

Development: It provides the tools to systematically explore and refine models.
CI/CD Integration: The artifacts and metadata from experiments feed directly into continuous integration and deployment pipelines. For example, a CI pipeline can automatically trigger when a new model is registered, running tests and preparing it for deployment.
Production Monitoring: The metrics and parameters from training runs serve as a baseline for monitoring the model's performance in production. If performance degrades (a concept known as model drift), the tracked experiments provide a clear, reproducible history to inform retraining and debugging efforts.

By maintaining a detailed record of every experiment, you create a transparent, auditable, and efficient workflow that accelerates the entire MLOps cycle.